import torch
import lightning.pytorch as pl
import numpy as np
from torch.utils.data import Dataset, DataLoader

class GreenFTrainDs(Dataset):
    def __init__(self, path, reaction=False):
        super().__init__()
        F = np.load(f'{path}/F.npy')
        if reaction:
            U = np.load(f'{path}/rU.npy')
            
        else:
            U = np.load(f'{path}/U.npy')
            
        self.U = torch.from_numpy(U).float()
        self.F = torch.from_numpy(F).float()
        self.length = F.shape[0] 
        
    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        u = self.U[idx]
        f = self.F[idx]
        data = f[None, ...]
        return  data, f[None, ...], u

class GreenFValDs(Dataset):
    def __init__(self, path, reaction=False):
        super().__init__()
        F = np.load(f'{path}/valF.npy')
        if reaction:
            U = np.load(f'{path}/valrU.npy')
        else:
            U = np.load(f'{path}/valU.npy')
        
        self.U = torch.from_numpy(U).float()
        self.F = torch.from_numpy(F).float()
        self.length = F.shape[0] 
        
    def __len__(self):
        return self.length
    
    def __getitem__(self, idx):
        u = self.U[idx]
        f = self.F[idx]
        data = f[None, ...]
        return  data, f[None, ...], u

class GreenRTrainDs(Dataset):
    def __init__(self, path, reaction=False):
        super().__init__()
        F = np.load(f'{path}/F.npy')
        R = np.load(f'{path}/R.npy')        
        if reaction:
            U = np.load(f'{path}/rU.npy')
            
        else:
            U = np.load(f'{path}/U.npy')
            
        self.U = torch.from_numpy(U).float()
        self.R = torch.from_numpy(R).float()
        self.F = torch.from_numpy(F).float()
        self.length = F.shape[0] 
        
    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        u = self.U[idx]
        r = self.R[idx]
        f = self.F[idx]
        # data = torch.stack([r, f], axis=0)
        return  r[None, ...], f[None, ...], u

class GreenRValDs(Dataset):
    def __init__(self, path, reaction=False):
        super().__init__()
        F = np.load(f'{path}/valF.npy')
        R = np.load(f'{path}/valR.npy')
        if reaction:
            U = np.load(f'{path}/valrU.npy')
        else:
            U = np.load(f'{path}/valU.npy')
    
        self.U = torch.from_numpy(U).float()
        self.R = torch.from_numpy(R).float()
        self.F = torch.from_numpy(F).float()
        self.length = F.shape[0] 
        
    def __len__(self):
        return self.length
    
    def __getitem__(self, idx):
        u = self.U[idx]
        r = self.R[idx]
        f = self.F[idx]
        # data = torch.stack([r, f], axis=0)
        return  r[None, ...], f[None, ...], u

class GreenRFTrainDs(Dataset):
    def __init__(self, path, reaction=False):
        super().__init__()
        F = np.load(f'{path}/F.npy')
        R = np.load(f'{path}/R.npy')        
        if reaction:
            U = np.load(f'{path}/rU.npy')
            
        else:
            U = np.load(f'{path}/U.npy')
            
        self.U = torch.from_numpy(U).float()
        self.R = torch.from_numpy(R).float()
        self.F = torch.from_numpy(F).float()
        self.length = F.shape[0] 
        
    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        u = self.U[idx]
        r = self.R[idx]
        f = self.F[idx]
        data = torch.stack([r, f], axis=0)
        return  data, f[None, ...], u

class GreenRFValDs(Dataset):
    def __init__(self, path, reaction=False):
        super().__init__()
        F = np.load(f'{path}/valF.npy')
        R = np.load(f'{path}/valR.npy')
        if reaction:
            U = np.load(f'{path}/valrU.npy')
        else:
            U = np.load(f'{path}/valU.npy')
    
        self.U = torch.from_numpy(U).float()
        self.R = torch.from_numpy(R).float()
        self.F = torch.from_numpy(F).float()
        self.length = F.shape[0] 
        
    def __len__(self):
        return self.length
    
    def __getitem__(self, idx):
        u = self.U[idx]
        r = self.R[idx]
        f = self.F[idx]
        data = torch.stack([r, f], axis=0)
        return  data, f[None, ...], u

class GreenXYFTrainDs(Dataset):
    def __init__(self, path, reaction=False, TrainN=2000):
        super().__init__()
        F = np.load(f'{path}/F.npy')
        if reaction:
            U = np.load(f'{path}/rU.npy')
            
        else:
            U = np.load(f'{path}/U.npy')

        self.TrainN = TrainN
        self.U = torch.from_numpy(U[:TrainN]).float()
        self.F = torch.from_numpy(F[:TrainN]).float()
        n = F.shape[-1]
        x = np.linspace(-1, 1, n)
        y = np.linspace(-1, 1, n) 
        xx, yy = np.meshgrid(x, y)
        self.xx, self.yy = torch.from_numpy(xx).float(), torch.from_numpy(yy).float()

    def __len__(self):
        return self.TrainN

    def __getitem__(self, idx):
        u = self.U[idx]
        f = self.F[idx]
        data = torch.stack([self.xx, self.yy, f], axis=0)
        return  data, f[None, ...], u

class GreenXYFValDs(Dataset):
    def __init__(self, path, reaction=False):
        super().__init__()
        F = np.load(f'{path}/valF.npy')
        if reaction:
            U = np.load(f'{path}/valrU.npy')
        else:
            U = np.load(f'{path}/valU.npy')
        
        self.U = torch.from_numpy(U).float()
        self.F = torch.from_numpy(F).float()
        self.length, n, _ = F.shape
        x = np.linspace(-1, 1, n)
        y = np.linspace(-1, 1, n) 
        xx, yy = np.meshgrid(x, y)
        self.xx, self.yy = torch.from_numpy(xx).float(), torch.from_numpy(yy).float()

    def __len__(self):
        return self.length
    
    def __getitem__(self, idx):
        u = self.U[idx]
        f = self.F[idx]
        data = torch.stack([self.xx, self.yy, f], axis=0)
        return  data, f[None, ...], u

class GreenDM(pl.LightningDataModule):
    
    def __init__(self, data_path, batch_size, n, data_type='xyf', reaction=False, TrainN=2000):
        super().__init__()
        self.path = f'{data_path}/{n}/Sample10000'
        self.batch_size = batch_size
        self.data_type = data_type
        self.tag = reaction
        self.TrainN = TrainN

    def setup(self, stage):    
        if stage == 'fit' or stage is None:
            match self.data_type:
                case 'xyf':
                    self.train_dataset = GreenXYFTrainDs(self.path, self.tag, self.TrainN)
                    self.val_dataset = GreenXYFValDs(self.path, self.tag)
                case 'rf':
                    self.train_dataset = GreenRFTrainDs(self.path, self.tag, self.TrainN)
                    self.val_dataset = GreenRFValDs(self.path, self.tag)
                case 'f':
                    self.train_dataset = GreenFTrainDs(self.path, self.tag, self.TrainN)
                    self.val_dataset = GreenFValDs(self.path)
                case 'r':
                    self.train_dataset = GreenRTrainDs(self.path, self.tag, self.TrainN)
                    self.val_dataset = GreenRValDs(self.path)
        else:
            pass

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=16)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=1, shuffle=False, num_workers=8)

if __name__ == '__main__':
    dm = GreenDM()